import os
import json
from PIL import Image
import shutil


# 定义类别和对应的group_id
category_to_group_id = {
    "black": 0,
    "blue": 1,
    "blue_safety_zone": 2,
    "blue_start": 3,
    "red": 4,
    "red_safety_zone": 5,
    "red_start": 6,
    "yellow": 7,
}

def move_images_to_target_folder(source_folder, target_folder):
    # 确保目标文件夹存在，如果不存在则创建
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)

    # 获取源文件夹中的所有图片文件
    image_files = [file for file in os.listdir(source_folder) if file.lower().endswith(('.png', '.jpg', '.jpeg'))]

    # 移动每个图片文件到目标文件夹
    for image_file in image_files:
        source_path = os.path.join(source_folder, image_file)
        target_path = os.path.join(target_folder, image_file)

        # 移动文件
        shutil.move(source_path, target_path)
        print(f"Moved {image_file} to {target_folder}")

def normalize(value, max_value):
    return value / max_value


def convert_to_yolov8_pose(json_file_path, output_file_path):
    with open(json_file_path, 'r') as file:
        data = json.load(file)

    # image_path = json_file_path.replace(".json", ".jpg")

    # 尝试打开图片文件，若文件不存在则跳过
    # try:
    #     image = Image.open(image_path)
    #     image_width, image_height = image.size
    # except FileNotFoundError:
    #     print(f"File not found: {image_path}. Skipping this file.")
    #     return

    image_width, image_height = data["imageWidth"], data["imageHeight"]
    yolo_format_data = []
    for shape in data["shapes"]:
        class_id = category_to_group_id.get(shape["label"])  # shape["label"] 类别名

        if shape["shape_type"] == "rectangle":
            center_x = (shape["points"][0][0] + shape["points"][1][0]) / 2
            center_y = (shape["points"][0][1] + shape["points"][1][1]) / 2
            width = abs(shape["points"][1][0] - shape["points"][0][0])
            height = abs(shape["points"][1][1] - shape["points"][0][1])

            # 归一化边界框坐标
            center_x = normalize(center_x, image_width)
            center_y = normalize(center_y, image_height)
            width = normalize(width, image_width)
            height = normalize(height, image_height)

            # 保留6位小数
            center_x = round(center_x, 6)
            center_y = round(center_y, 6)
            width = round(width, 6)
            height = round(height, 6)

            yolo_format_data.append(f"{class_id} {center_x} {center_y} {width} {height} ")

    with open(output_file_path, 'w') as output_file:
        for line in yolo_format_data:
            output_file.write(line + "\n")

# 指定存储标注文件和图像文件的文件夹路径
folder_path = "/home/champrin/Desktop/校赛小球数据集/目标识别/Engineering training 2.v1i.yolov8/json_annotations"
output_folder = os.path.join(os.path.dirname(folder_path), "labels")

os.makedirs(output_folder, exist_ok=True)

for filename in os.listdir(folder_path):
    if filename.endswith(".json"):
        json_file_path = os.path.join(folder_path, filename)
        output_file_path = os.path.join(output_folder, filename.replace(".json", ".txt"))

        convert_to_yolov8_pose(json_file_path, output_file_path)

print("Conversion completed.")

# 执行移动操作
move_images_to_target_folder(folder_path, output_folder)